package com.gitee.liuzhi.common.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.BatchResult;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.stereotype.Component;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import javax.annotation.Resource;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiFunction;

/**
 * Created by 2023/8/31/0031
 *
 * @author LiuZhi
 * @version V1.0
 * @description: 使用 Mybatis-Executor 的 Batch 模式批量执行, 适用于批量 Update、Insert、Delete
 */
@Component
@Slf4j
public class BatchUpsert {

    /*
    JDBC 在执行 SQL 语句时, 会将 SQL 语句以及实参通过网络请求的方式发送到数据库, 一次执行一条 SQL 语句, 一方面会减小请求包的有效负载, 另一个方面会增加耗费在网络通信上的时间
    org.apache.ibatis.session.ExecutorType.BATCH 批处理的处理方式
    + 可以在 JDBC 客户端缓存多条 SQL语句, 然后在 flush 或缓存满的时候, 将多条 SQL 语句打包发送到数据库执行, 这样就可以有效地降低上述两方面的损耗, 从而提高系统性能
    + 在使用 ExecutorType.BATCH 时, MyBatis 不会立即执行 SQL, 而是将 SQL 缓存起来, 等到调用 SqlSession.commit() 或 SqlSession.flushStatements() 时才会一次性执行
    + 如果在批处理过程中出现异常, 需要确保进行适当的错误处理, 例如回滚事务
    + 在使用 ExecutorType.BATCH 时, insert、update 和 delete 操作返回的结果可能不正确 (通常为 -2147482646), 如果需要获取这些操作的结果, 可以在 SqlSession.flushStatements() 之后调用 SqlSession.clearCache()，然后再执行查询操作
     */

    private static final int BATCH = 1000;

    @Resource
    private SqlSessionFactory sqlSessionFactory;

    /**
     * 批处理方式执行
     *
     * @param data        数据集合
     * @param mapperClass Mapper的Class对象
     * @param function    Mapper的执行方法
     * @return 操作的行数
     */
    public <T, U, R> int execute(List<T> data, Class<U> mapperClass, BiFunction<T, U, R> function, int pageSize) {
        int i = 1;
        SqlSession batchSqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
        AtomicInteger totalRowsUpdated = new AtomicInteger();
        try {
            U mapper = batchSqlSession.getMapper(mapperClass);
            int size = data.size();
            for (T element : data) {
                function.apply(element, mapper);
                if ((i % pageSize == 0) || i == size) {
                    List<BatchResult> batchResult = batchSqlSession.flushStatements();
                    batchResult.forEach(e -> {
                        for (int updateCount : e.getUpdateCounts()) {
                            // 累加更新行数
                            totalRowsUpdated.addAndGet(updateCount);
                        }
                    });
                }
                i++;
            }
            // 非事务环境下强制 commit，事务情况下该 commit 相当于无效 (如果存在外部事务,则和外部事务一起提交,避免异常时该方式不回滚)
            batchSqlSession.commit(!TransactionSynchronizationManager.isSynchronizationActive());
        } catch (Exception e) {
            batchSqlSession.rollback();
            log.error(String.format("BatchUpsert execute fail. message : %s", e.getMessage()), e);
            throw e;
        } finally {
            batchSqlSession.close();
        }
        return totalRowsUpdated.get();
    }


    /**
     * 批处理方式执行
     * <p>
     * + batchUpsert.execute(commitList, WmsLocationStockMapper.class, (item, locationStockMapper) -> locationStockMapper.updateCommitStock(item));
     *
     * @param data        数据集合
     * @param mapperClass Mapper的Class对象
     * @param function    Mapper的执行方法 (自定义处理逻辑)
     * @param <T>         数据类型
     * @param <U>         Mapper的类型
     * @return 受影响行数 (注意: 在批量insert时可能出现插入数量值不准确)
     */
    public <T, U, R> int execute(List<T> data, Class<U> mapperClass, BiFunction<T, U, R> function) {
        return execute(data, mapperClass, function, BATCH);
    }

}
